# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
束搜索（Beam Search）算法实现

本模块实现了基于束搜索的推理算法，用于解决数学推理问题。
束搜索是一种启发式搜索算法，在每一步保留最有前途的 k 个候选路径，
通过贪心策略在搜索空间中寻找最优解。

主要功能:
- 实现步级束搜索算法的核心流程
- 支持多候选路径的并行扩展
- 提供价值估计和节点排序机制
- 集成代码执行和错误处理
- 支持最终答案的收集和排序

核心组件:
- BS: 主要的束搜索算法实现类
- 节点管理: 维护当前节点、候选节点和最终答案节点
- 价值评估: 基于价值估计对候选路径进行排序
- 路径扩展: 根据 LLM 输出创建新的候选路径
"""
# 导入必要的库和模块
from __future__ import annotations  # 启用延迟注解，支持前向引用类型提示
from termcolor import colored  # 终端彩色输出库，用于调试信息的可视化显示
from typing import Dict, Any, Optional, Type, List, Tuple, Callable, Union  # 类型提示工具，提供静态类型检查支持
from pydantic import BaseModel, PrivateAttr, conlist, ConfigDict, field_validator  # Pydantic 数据验证和设置管理
from functools import partial  # 函数工具，用于创建偏函数
from vllm.outputs import RequestOutput  # VLLM 输出类型，处理语言模型的生成结果
from rstar_deepthink.nodes.base_node import BaseNode  # 搜索树节点基类
from rstar_deepthink.constants import (  # 项目常量定义
    TOO_MANY_CODE_ERRORS,  # 连续代码错误过多的终止标记
    TOO_MANY_STEPS,       # 步数超限的终止标记
    NO_VALID_CHILD,       # 无有效子节点的终止标记
    CODE_END,             # 代码结束标记
    ANSWER,               # 答案标记
    OUTPUT,               # 输出标记
    OUTPUT_END,           # 输出结束标记
)
from .tree import BaseTree, code_execution  # 搜索树基类和代码执行函数


class BS(BaseTree):
    """
    步级束搜索（Step-level Beam Search）算法实现类
    
    继承自 BaseTree 类，实现了完整的束搜索算法流程。
    束搜索通过在每一步保留最有前途的 k 个候选路径来平衡搜索效率和解的质量。
    与 MCTS 不同，束搜索采用贪心策略，不进行回溯和价值传播。
    
    算法特点:
    1. 贪心搜索：每步只保留价值最高的 k 个候选
    2. 并行扩展：同时扩展多个候选路径
    3. 价值排序：基于价值估计对候选进行排序
    4. 终止条件：达到最大深度或找到有效答案时停止
    
    属性:
        NODE_KEYS: 节点状态字典的键列表
        prompt_wrap: 提示词包装函数
        obs_wrap: 观察结果包装函数
        step_unwrap: 步骤结果解析函数
        current_top_num: 当前保留的候选数量
        current_nodes: 当前活跃的节点列表
        final_answer_nodes: 包含最终答案的节点列表
        candidate_nodes: 候选节点列表
        rollout_idx: 当前推演轮次索引
    
    核心方法:
    - select_next_step: 选择下一步要扩展的节点
    - generate_next_step: 生成下一步的候选节点
    - create_child: 创建子节点
    - get_steps: 获取排序后的解决方案
    """
    # 节点状态字典的必需键列表，定义了每个节点必须包含的状态信息
    NODE_KEYS: List[str] = ["action", "action_input", "final_answer"]
    
    # 函数接口定义，用于处理不同格式的输入输出
    prompt_wrap: Optional[Callable[[...], str]] = None      # 提示词包装函数，将问题和部分解转换为模型输入
    obs_wrap: Optional[Callable[str, str]] = None           # 观察结果包装函数，格式化代码执行结果
    step_unwrap: Optional[Callable[[...], Dict[str, str]]] = None  # 步骤结果解析函数，解析模型输出
    
    # 束搜索算法的核心参数和状态
    current_top_num: int = 1                                # 当前保留的最优候选数量（束宽度）
    current_nodes: List[Type[BaseNode]] = []                # 当前活跃的节点列表，等待扩展
    final_answer_nodes: List[Type[BaseNode]] = []           # 包含有效最终答案的节点列表
    candidate_nodes: List[Type[BaseNode]] = []              # 候选节点列表，用于价值评估和选择
    rollout_idx: int = 0                                    # 当前推演轮次的索引
    def __init__(self, **kwargs) -> None:
        """
        初始化束搜索算法实例
        
        设置必要的函数接口和初始状态，根据配置选择合适的提示词格式。
        
        参数:
            **kwargs: 传递给父类的关键字参数
            
        初始化流程:
        1. 调用父类初始化方法
        2. 根据配置设置提示词处理函数
        3. 将根节点添加到候选列表
        4. 设置束宽度参数
        """
        super().__init__(**kwargs)

        # 根据配置选择提示词格式和处理函数
        if self.config.prompt_wrap == "rstar":
            from .utils import rstar_prompt_wrap, rstar_obs_wrap, rstar_step_result_unwrap

            self.prompt_wrap = rstar_prompt_wrap      # 设置 R* 格式的提示词包装函数
            self.obs_wrap = rstar_obs_wrap            # 设置观察结果格式化函数
            self.step_unwrap = rstar_step_result_unwrap  # 设置步骤结果解析函数

        # 初始化搜索状态
        self.candidate_nodes.append(self.current_node)  # 将根节点作为第一个候选
        self.current_top_num = self.config.step_beam_width  # 设置束宽度

    @field_validator("config")
    def validate_config(cls, cfg: Any):
        """
        验证束搜索配置参数
        
        确保配置参数符合束搜索算法的要求，包括生成样本数和停止条件。
        
        参数:
            cfg: 配置对象
            
        返回:
            验证后的配置对象
            
        异常:
            ValueError: 当配置参数不符合要求时抛出
        """
        BaseTree.validate_config(cfg)
        if not cfg.n_generate_sample >= 1:
            raise ValueError(f"Wrong value for config n_generate_sample, must be greater than 1")
        if cfg.stop is None:
            raise ValueError(f"Wrong value for config stop, cannot be None")
        return cfg
    

    def is_ignored_node(self, node: Type[BaseNode]) -> bool:
        """
        判断节点是否应该被忽略
        
        检查节点是否已经终止或超过最大深度限制。
        被忽略的节点不会参与后续的扩展过程。
        
        参数:
            node: 要检查的节点
            
        返回:
            bool: 如果节点应该被忽略则返回 True，否则返回 False
        """
        return node.is_terminal or node.depth > self.config.max_depth

    def should_generate_next(self) -> bool:
        """
        判断是否需要生成下一步
        
        检查当前节点列表中是否还有可以继续扩展的节点。
        如果所有节点都已终止或超过深度限制，则不需要继续生成。
        
        返回:
            bool: 如果需要生成下一步则返回 True，否则返回 False
        """
        need_generate = False
        for step_node in self.current_nodes:
            if not self.is_ignored_node(step_node):
                need_generate = True
                break
        return need_generate
    
    def has_expanded(self) -> bool:
        """
        检查当前节点是否已经扩展
        
        判断当前节点列表中的第一个节点是否已经有子节点。
        用于确定搜索树的扩展状态。
        
        返回:
            bool: 如果已经扩展则返回 True，否则返回 False
        """
        if not self.current_nodes:
            return False
        step_node = self.current_nodes[0]
        if step_node.has_children():
            return True
        return False

    def get_rewards(self):
        """
        获取当前节点的奖励值列表
        
        收集当前所有活跃节点的奖励值，用于评估和调试。
        如果节点没有奖励值，则使用默认值 0。
        
        返回:
            List[float]: 当前节点的奖励值列表
        """
        rewards = []
        for node in self.current_nodes:
            rewards.append(node.reward if node.reward is not None else 0)  # 默认奖励为 0
        return rewards

    def create_prompt(
        self,
        is_value_only: bool = False,
    ) -> str:
        """
        if is_value_only, the prompt is used to produce value estimate.
        """
        prompts = []
        current_nodes = self.candidate_nodes if is_value_only else self.current_nodes
        for current_node in current_nodes:
            if not is_value_only and self.is_ignored_node(current_node):
                continue
            partial_solution = self.collect_partial_solution(current_node)
            prompt = self.prompt_wrap(
                self.question, 
                partial_solution,
                self.config
            )
            if is_value_only:
                prompt = {
                    "prefix": "",
                    "text": prompt,
                }
            prompts.append(prompt)
        return prompts
    
    @staticmethod
    def is_valid_final_answer_node(node: Type[BaseNode]) -> bool:
        if node.is_terminal and node.state["final_answer"] and \
           node.state["final_answer"] not in [NO_VALID_CHILD, TOO_MANY_STEPS, TOO_MANY_CODE_ERRORS]:
            return True
        return False
    
    def select_next_step(self, outputs=None, from_root=False) -> None:
        """
        选择下一步要扩展的节点
        
        这是束搜索算法的核心选择机制，负责根据价值估计选择最有前途的候选节点。
        该方法处理价值更新、节点排序、最终答案收集和束宽度控制。
        
        参数:
            outputs: 价值估计输出结果，用于更新候选节点的价值
            from_root: 是否从根节点开始选择（当前未使用）
            
        算法流程:
        1. 清空当前节点列表
        2. 如果有价值估计输出，更新候选节点的价值
        3. 按价值对候选节点进行降序排序
        4. 识别并收集包含有效最终答案的节点
        5. 移除终止节点和超深度节点
        6. 根据束宽度选择最优的前 k 个节点
        
        注意:
        - 有效的最终答案节点会被添加到 final_answer_nodes 列表
        - 束宽度会根据找到的最终答案数量动态调整
        """
        self.current_nodes = []
        if outputs is not None:
            for candidate_node, output in zip(self.candidate_nodes, outputs):
                candidate_node.value = output.value_estimate if output.value_estimate is not None else 0
            
        self.candidate_nodes = sorted(self.candidate_nodes, key=lambda x: x.value, reverse=True)
        self.current_nodes = self.candidate_nodes[:]

        for current_node in self.current_nodes[:]: 
            if self.__class__.is_valid_final_answer_node(current_node):
                self.final_answer_nodes.append(current_node)
                self.current_nodes.remove(current_node)
                self.current_top_num -= 1
            elif current_node.is_terminal or current_node.depth > self.config.max_depth:
                self.current_nodes.remove(current_node)
        self.current_nodes = self.candidate_nodes[:self.current_top_num]
        
    def generate_next_step(self, outputs: List[RequestOutput]) -> None:
        """
        从 LLM 输出生成下一步候选节点
        
        这是束搜索算法的核心扩展方法，负责处理语言模型的输出并创建新的候选节点。
        该方法解析 LLM 生成的文本，创建子节点，并将它们添加到候选列表中。
        
        参数:
            outputs: 语言模型的输出结果列表，每个输出对应一个当前节点
            
        算法流程:
        1. 清空候选节点列表
        2. 遍历每个当前节点和对应的 LLM 输出
        3. 处理每个输出的多个生成候选
        4. 解析生成文本为结构化的步骤结果
        5. 为每个有效结果创建子节点
        6. 将所有新创建的子节点添加到候选列表
        
        注意:
        - 每个节点可能产生多个子节点（基于 n_generate_sample 配置）
        - 停止原因会被附加到生成文本中以确保完整性
        - 子节点的创建包括状态更新、代码执行和终止条件检查
        """
        self.candidate_nodes = []
        for current_node, output in zip(self.current_nodes, outputs):
            self.current_node = current_node
            for idx, output in enumerate(output.outputs):
                if not output.stop_reason: output.stop_reason = ""
                step_result, parser_result = self.step_unwrap(output.text + output.stop_reason)
                self.create_child(step_result, parser_result, current_node)
            self.candidate_nodes.extend(current_node.children)

    def create_node(self, parent: Optional[Type[BaseNode]] = None) -> Type[BaseNode]:
        return BaseNode(
            parent=parent, 
            additional_state_keys=self.NODE_KEYS,
        )
    
    def create_child(
        self, 
        step_result: str, 
        parser_result: Dict[str, str], 
        node: Type[BaseNode],
    ) -> None:
        """
        从步骤结果创建子节点
        
        根据解析的步骤结果创建新的子节点，处理不同类型的动作和终止条件。
        这个方法负责维护搜索树的结构和节点状态的一致性。
        
        参数:
            step_result: 原始的步骤结果文本
            parser_result: 解析后的结构化结果字典
            node: 父节点
            
        处理逻辑:
        1. 创建子节点并设置基本属性（标签、深度）
        2. 根据解析结果的类型进行不同处理：
           - 解析失败：标记为无效子节点
           - 最终答案：标记为终止节点
           - 动作执行：执行代码并处理观察结果
           - 普通文本：直接设置文本状态
        3. 检查深度限制和错误条件
        4. 设置适当的终止状态和最终答案
        
        错误处理:
        - 代码执行错误会导致节点终止
        - 超过最大深度会强制终止
        - 连续错误过多会标记为 TOO_MANY_CODE_ERRORS
        """
        new_node = self.create_node(parent=node)
        parent_child_count = len(node.children)
        new_node.tag = f"{node.tag}.{parent_child_count + 1}"
        new_node.depth = node.depth + 1

        if parser_result is None:
            new_node.is_terminal = True
            new_node.state["text"] = step_result
            new_node.state["final_answer"] = NO_VALID_CHILD
        elif parser_result["final_answer"]:
            new_node.is_terminal = True
            new_node.state["text"] = step_result
            new_node.state["final_answer"] = parser_result["final_answer"]
        elif parser_result["action"]:
            observation = code_execution(node, parser_result)
            new_node.state["action"] = parser_result["action"]
            new_node.state["action_input"] = parser_result["action_input"]
            new_node.state["observation"] = observation
            if CODE_END in parser_result["action_input"]:
                observation = self.obs_wrap(observation)
                new_node.state["text"] = f"{step_result}{self.config.step_delim}{observation}"
                if "Error" in observation:
                    new_node.is_terminal = True
                    new_node.state["final_answer"] = TOO_MANY_CODE_ERRORS
            else:
                new_node.state["text"] = step_result
                
            if "error" in observation.lower():
                observation = self.obs_wrap(observation)
                step_result = step_result + CODE_END if CODE_END not in step_result else step_result
                new_node.state["text"] = f"{step_result}{self.config.step_delim}{observation}"
                new_node.is_terminal = True
                new_node.state["final_answer"] = TOO_MANY_CODE_ERRORS

        else:
            new_node.state["text"] = step_result

        if not new_node.is_terminal and new_node.depth > self.config.max_depth:
            new_node.is_terminal = True
            new_node.state["final_answer"] = TOO_MANY_STEPS

        node.children.append(new_node)

    def get_steps(self):
        """
        获取最终的解决方案步骤及其价值
        
        收集并排序所有包含最终答案的节点，构建完整的解决方案信息。
        每个解决方案包含问题、答案、价值评分和完整的求解路径。
        
        返回:
            List[Dict]: 解决方案列表，每个解决方案包含以下字段：
                - question: 原始问题
                - ground_truth: 标准答案（如果有）
                - value: 节点的价值评分
                - final_answer: 最终答案
                - solution: 完整的求解路径
                - tag: 节点标识符
                
        算法流程:
        1. 遍历所有包含最终答案的节点
        2. 为每个节点构建完整的解决方案信息
        3. 收集求解路径（从根到叶的完整步骤）
        4. 按价值对解决方案进行降序排序
        
        注意:
        - 只处理已经找到最终答案的节点
        - 价值评分用于排序，优先返回最有前途的解决方案
        - 解决方案路径通过 collect_partial_solution 方法构建
        """
        final_answer_states = []
        for cur_node in self.final_answer_nodes:
            states = {
                "question": self.question,
                "ground_truth": self.ground_truth,
                "value": cur_node.value,
                "final_answer": cur_node.state["final_answer"],
                "solution": self.collect_partial_solution(cur_node),
                "tag": cur_node.tag,
            }
            final_answer_states.append(states)

        solutions = sorted(final_answer_states, key=lambda x: x['value'], reverse=True)
        return solutions

    def return_states(self) -> Dict[str, Union[Any, Dict[str, str]]]:
        """
        返回搜索树的完整状态信息
        
        遍历整个搜索树，收集所有节点的状态信息和最终解决方案。
        这个方法主要用于调试、分析和结果展示。
        
        返回:
            Dict[str, Union[Any, Dict[str, str]]]: 包含以下信息的字典：
                - 每个节点的标签作为键，节点状态作为值
                - 每个节点状态包含其价值评分
                - "solutions" 键包含所有最终解决方案
                
        遍历算法:
        1. 使用广度优先搜索遍历整个搜索树
        2. 收集每个节点的完整状态信息
        3. 添加节点的价值评分到状态中
        4. 将所有最终解决方案添加到结果中
        
        用途:
        - 调试搜索过程
        - 分析节点扩展情况
        - 可视化搜索树结构
        - 导出完整的搜索历史
        """
        candidates = [self.root]
        states = {}
        while candidates:
            node = candidates.pop(0)
            states[node.tag] = node.state
            states[node.tag]["value"] = node.value
            if node.has_children():
                candidates.extend(node.children)
        states["solutions"] = self.get_steps()
        return states
